#include "CodecNvJPEG2000.h"

#include "cuda_runtime.h"
#include "device_launch_parameters.h"
#include "nvjpeg2k.h"

#ifndef QT_VERSION
#include <iostream>
#endif
#include <float.h>

class JP2_Decoder::D
{
public:
    bool isvalid = true;
    int cudaDevNum = 0;
    cudaStream_t cudastream = nullptr;
    nvjpeg2kHandle_t nvjpeg2k_handle = nullptr;
    nvjpeg2kStream_t nvjpeg2k_stream = nullptr;
    nvjpeg2kDecodeState_t decode_state = nullptr;

    const unsigned char * input_buf = nullptr;  //  经过编码后的图像数据
    size_t input_buf_len = 0;  //  经过编码后的图像数据字节长度
    nvjpeg2kImageInfo_t image_info;   //  解码后图像信息
    nvjpeg2kImageComponentInfo_t image_comp_info[4];  //  解码后通道信息， 应该不会超过4通道吧

    char * output_buf[4] = {nullptr};   //  输出缓冲区
    size_t output_buf_pitch[4] = {0};   //  输出缓冲区每行特殊对齐填充后的长度
    size_t output_buf_len[4] = {0};     //  输出缓冲区字节数
    nvjpeg2kImage_t output_image;

    char * output_buf_host[4] = {nullptr};   //  cpu 输出缓冲区
};

JP2_Decoder::JP2_Decoder(int cudaDevNum):
    m_dptr(new D)
{
    std::memset(&m_dptr->output_image,0,sizeof (nvjpeg2kImage_t));
    std::memset(&m_dptr->image_info,0,sizeof (nvjpeg2kImageInfo_t));
    std::memset(&m_dptr->image_comp_info[0],0,
            4*sizeof (nvjpeg2kImageComponentInfo_t));

    m_dptr->cudaDevNum = cudaDevNum;
    int count = 0;
    int driverVersion = 0;
    cudaError_t ret = cudaGetDeviceCount(& count);
    cudaDriverGetVersion(&driverVersion);
    m_dptr->isvalid = m_dptr->isvalid && ret==cudaSuccess;
    m_dptr->isvalid = m_dptr->isvalid && count > 0;
    m_dptr->isvalid = m_dptr->isvalid && m_dptr->cudaDevNum < count;
    m_dptr->isvalid = m_dptr->isvalid && driverVersion>=1100;

    if(m_dptr->isvalid){
        ret = cudaSetDevice(m_dptr->cudaDevNum);
        m_dptr->isvalid = m_dptr->isvalid && ret==cudaSuccess;
        auto j2k_ret = nvjpeg2kCreateSimple(&m_dptr->nvjpeg2k_handle);
        m_dptr->isvalid = m_dptr->isvalid && NVJPEG2K_STATUS_SUCCESS == j2k_ret;
        nvjpeg2kDecodeStateCreate(m_dptr->nvjpeg2k_handle,&m_dptr->decode_state);
        m_dptr->isvalid = m_dptr->isvalid && NVJPEG2K_STATUS_SUCCESS == j2k_ret;
        nvjpeg2kStreamCreate(&m_dptr->nvjpeg2k_stream);
        m_dptr->isvalid = m_dptr->isvalid && NVJPEG2K_STATUS_SUCCESS == j2k_ret;
        m_dptr->output_image.pixel_data = reinterpret_cast<void **>(m_dptr->output_buf);
        m_dptr->output_image.pitch_in_bytes = m_dptr->output_buf_pitch;

        ret = cudaStreamCreate(&m_dptr->cudastream);
        m_dptr->isvalid = m_dptr->isvalid && ret==cudaSuccess;

        m_dptr->isvalid = true;
    }
}

JP2_Decoder::~JP2_Decoder()
{
    if(m_dptr->isvalid){
        cudaStreamDestroy(m_dptr->cudastream);
    }

    if(m_dptr->nvjpeg2k_stream){
        nvjpeg2kStreamDestroy(m_dptr->nvjpeg2k_stream);
        m_dptr->nvjpeg2k_stream = nullptr;
    }

    if(m_dptr->decode_state){
        nvjpeg2kDecodeStateDestroy(m_dptr->decode_state);
        m_dptr->decode_state = nullptr;
    }

    if(m_dptr->nvjpeg2k_handle){
        nvjpeg2kDestroy(m_dptr->nvjpeg2k_handle);
        m_dptr->nvjpeg2k_handle = nullptr;
    }

    for (int i = 0; i < 4; ++i) {
        if(m_dptr->output_buf[i]){
            cudaFree(m_dptr->output_buf[i]);
            m_dptr->output_buf[i] = nullptr;
        }

        if(m_dptr->output_buf_host[i]){
            delete m_dptr->output_buf_host[i];
            m_dptr->output_buf_host[i] = nullptr;
        }
    }

    delete m_dptr;
}

bool JP2_Decoder::isValid()
{
    return m_dptr->isvalid;
}

bool JP2_Decoder::parseEncodedBuf(const char *buf, size_t len)
{
    if(!buf || len<1) return false;
    if(!m_dptr->isvalid) return false;

    auto j2k_ret = nvjpeg2kStreamParse(m_dptr->nvjpeg2k_handle,
                                       reinterpret_cast<const unsigned char *>(buf), len,
                                       0, 0,
                                       m_dptr->nvjpeg2k_stream);

    if(j2k_ret!=NVJPEG2K_STATUS_SUCCESS) return false;

    m_dptr->input_buf = reinterpret_cast<const unsigned char *>(buf);
    m_dptr->input_buf_len = len;

    j2k_ret = nvjpeg2kStreamGetImageInfo(m_dptr->nvjpeg2k_stream,
                                         &m_dptr->image_info);
    if(j2k_ret!=NVJPEG2K_STATUS_SUCCESS) return false;

    int c = static_cast<int>(m_dptr->image_info.num_components);
    if(c > 4) c=4;
    if(c < 1) c=1;
    for (int i = 0; i < c; ++i) {
        j2k_ret = nvjpeg2kStreamGetImageComponentInfo(
                    m_dptr->nvjpeg2k_stream,
                    &m_dptr->image_comp_info[i],
                    static_cast<unsigned int>(i)
                    );
        if(j2k_ret!=NVJPEG2K_STATUS_SUCCESS) return false;
    }

    cudaError_t cuda_ret = cudaSuccess;
    for (int i = 0; i < c; ++i) {
        size_t bytepercomponent = (m_dptr->image_comp_info[i].precision + 7) / 8;
        size_t buflen = m_dptr->image_comp_info[i].component_height *
                m_dptr->image_comp_info[i].component_width *
                bytepercomponent;
        if(buflen > m_dptr->output_buf_len[i]){
            if(m_dptr->output_buf[i]) cudaFree(m_dptr->output_buf[i]);
            if(m_dptr->output_buf_host[i]) delete m_dptr->output_buf_host[i];
            m_dptr->output_buf[i] = nullptr;
            m_dptr->output_buf_host[i] = nullptr;
        }
        if(!m_dptr->output_buf[i]){
            cuda_ret = cudaMallocPitch(&m_dptr->output_buf[i],&m_dptr->output_buf_pitch[i],
                                       m_dptr->image_comp_info[i].component_width * bytepercomponent,
                                       m_dptr->image_comp_info[i].component_height
                                       );
            if(cuda_ret != cudaSuccess){
                m_dptr->output_buf[i] = nullptr;
                return false;
            }
        }
        m_dptr->output_buf_len[i] = buflen;
    }

    for (int i = 0; i < c; ++i) {
        if(!m_dptr->output_buf_host[i]){
            m_dptr->output_buf_host[i] = new char[m_dptr->output_buf_len[i]];
        }
    }

    m_dptr->output_image.pixel_type = nvjpeg2kImageType_t::NVJPEG2K_UINT8;
    m_dptr->output_image.num_components = m_dptr->image_info.num_components;
    if(m_dptr->image_comp_info[0].precision > 8){
        m_dptr->output_image.pixel_type = nvjpeg2kImageType_t::NVJPEG2K_UINT16;
    }

    return true;
}

bool JP2_Decoder::isU16C1Image()
{
    if(!m_dptr->isvalid) return false;
    bool ret = m_dptr->image_info.num_components == 1;
    ret = ret && m_dptr->image_comp_info[0].precision == 16;
    ret = ret && m_dptr->image_comp_info[0].sgn == 0;
    return ret;
}

void JP2_Decoder::getOutputInfo(int &width, int &height, int component)
{
    width = 0;
    height = 0;

    if(component < 0) return;

    if(static_cast<unsigned int>(component) >= m_dptr->image_info.num_components)
        return;

    if(component > 3) return;

    width = static_cast<int>(m_dptr->image_comp_info[component].component_width);
    height = static_cast<int>(m_dptr->image_comp_info[component].component_height);

}

bool JP2_Decoder::doDecode()
{
    if(!m_dptr->isvalid) return false;

    auto j2k_ret = nvjpeg2kDecode(m_dptr->nvjpeg2k_handle,
                                  m_dptr->decode_state,
                                  m_dptr->nvjpeg2k_stream,
                                  &m_dptr->output_image,
                                  m_dptr->cudastream);
    if(j2k_ret!=NVJPEG2K_STATUS_SUCCESS) return false;

    cudaError_t cuda_ret = cudaSuccess;
    for (int i = 0; i < static_cast<int>(m_dptr->output_image.num_components); ++i) {
        size_t pixbyte = 1;
        if(m_dptr->output_image.pixel_type == NVJPEG2K_UINT16 ){
            pixbyte = 2;
        }
        cudaMemcpy2DAsync(m_dptr->output_buf_host[i],
                          m_dptr->image_comp_info[i].component_width * pixbyte,
                          m_dptr->output_buf[i],
                          m_dptr->output_buf_pitch[i],
                          m_dptr->image_comp_info[i].component_width * pixbyte,
                          m_dptr->image_comp_info[i].component_height,
                          cudaMemcpyDeviceToHost,
                          m_dptr->cudastream
                          );
    }

    cuda_ret = cudaStreamSynchronize(m_dptr->cudastream);
    if(cuda_ret!=cudaSuccess) return false;

    return true;
}

const char *JP2_Decoder::getDecodedRes(size_t &ressize,int component)
{
   if(!m_dptr->isvalid) return nullptr;
   if(component<0 || component>3) return nullptr;

   ressize = m_dptr->output_buf_len[component];
   return m_dptr->output_buf_host[component];
}
